import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.utils import resample
import matplotlib as mpl
from matplotlib import font_manager as fm
from utils.data_split import DataSplitUtil
from utils.utils import *
# 指定 Times New Roman 字体的路径
font_path = '/usr/share/fonts/truetype/msttcorefonts/Times_New_Roman.ttf'
import joblib
# 创建 FontProperties 对象
times_new_roman_prop = fm.FontProperties(fname=font_path)

# 注册字体到 matplotlib
mpl.rcParams['font.family'] = times_new_roman_prop.get_name()

# 设置负号显示正常（防止负号变成方块）
mpl.rcParams['axes.unicode_minus'] = False
from config.config import *
from utils.ensemble import AverageEnsemble  # 确保 utils/ensemble.py 中定义了 AverageEnsemble

# =======================================
# Step 1. 数据读取与特征选择
# =======================================

df = pd.read_csv(signature_score_csv_path)
# cohort_list = ['train', 'test','all']
cohort_list = [ 'test']
for cohort in cohort_list:
    train_df , test_df = DataSplitUtil(split_random_state_list[0]).get_train_test_df(df)
    if cohort == 'train':
        y = train_df[target_column]
        X = train_df.drop(columns=exclude_columns)
    elif cohort == 'test':
        y = test_df[target_column]
        X = test_df.drop(columns=exclude_columns)
    else:
        X = df.drop(columns=exclude_columns)
        y = df[target_column]
    N=X.shape[0]
    save_path = opj(shap_result_path,cohort)
    md(save_path)

    # =========================
    # 2. 计算SHAP值，并可视化
    # =========================
    import shap

    # 假设你已有：
    # X: 包含 6 个特征的 DataFrame，shape=(n_samples, 6)
    # ensemble: 训练好的 AverageEnsemble 对象

    # 1. 读取保存的 Voting 模型
    import joblib
    load_path = opj(ensemble_model_weight_path,'model.joblib')
    ensemble = joblib.load(load_path)
    print(f"Voting Ensemble 模型已加载：{load_path}")

    # 1. 背景样本：这里直接用全部 X（6 个特征，样本量一般也不会太大）
    background = X.values

    # 2. 定义 ensemble 的预测函数
    def ensemble_predict_proba(data):
        # data: numpy array 或者 DataFrame
        df = pd.DataFrame(data, columns=X.columns) if not isinstance(data, pd.DataFrame) else data
        return ensemble.predict_proba_by_dif(df)[:, 1]

    # 3. 创建 KernelExplainer（对 6 维输入，用 logit link）
    explainer = shap.KernelExplainer(
        model=ensemble_predict_proba,
        data=background,
        link="logit"
    )

    pkl_path = opj(ensemble_model_weight_path,f'shap_{cohort}.pkl')
    if not ope(pkl_path):
        # 4. 计算 SHAP 值
        # shap_value不存在
        shap_values = explainer.shap_values(background)

        # 保存 shap_values
        joblib.dump(shap_values,opj(ensemble_model_weight_path,f'shap_{cohort}.pkl'))
        print(f"{cohort} 队列的SHAP 值已保存")

    # 加载 shap_values
    shap_values = joblib.load(opj(ensemble_model_weight_path,f'shap_{cohort}.pkl'))
    print(f"{cohort} 队列的 SHAP 值已加载")

    X_plot = X.copy()
    X_plot.columns = [
        col.replace("Factor ", "Phenotype ").replace(" Score", "") if "Factor" in col else col
        for col in X_plot.columns]

    if cohort != 'all':
        mean_abs_shap = np.abs(shap_values).mean(axis=0)
        feature_names = X_plot.columns

        sort_idx = np.argsort(mean_abs_shap)[::-1]
        sorted_values = mean_abs_shap[sort_idx]
        sorted_names = feature_names[sort_idx]

        # 归一化相对重要性
        relative_importance = sorted_values / sorted_values.sum()

        # 生成颜色（比如 Purples 颜色映射）
        custom_colors = ['#1E88E5','#099B57', '#FF0D57', '#FF9F3F',]

        # 画条形图
        fig, ax = plt.subplots(figsize=(5.5, 4))

        bars = ax.barh(
            y=np.arange(len(sorted_names)),
            width=sorted_values,
            color=custom_colors,
            height=0.6,
        )
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)

        ax.set_yticks(np.arange(len(sorted_names)))
        ax.set_yticklabels(sorted_names, fontsize=14)
        ax.set_xlabel("mean(|SHAP value|)", fontsize=14)
        ax.tick_params(axis='x', labelsize=14)
        ax.invert_yaxis()  # 最大值在上
        ax.set_xlim(left=0)

        # 添加右下角饼图
        from mpl_toolkits.axes_grid1.inset_locator import inset_axes
        ax_inset = inset_axes(ax, width="32%", height="32%", loc='lower right', bbox_to_anchor=(0.20, 0, 1, 1), bbox_transform=ax.transAxes, borderpad=1.0)
        wedges, texts, autotexts = ax_inset.pie(
            relative_importance,
            labels=None,
            colors=custom_colors,
            startangle=140,
            autopct='%1.1f%%',
            pctdistance=1.5,
            textprops={'fontsize': 10},
            wedgeprops=dict(width=0.5, edgecolor='white')
        )

        for autotext in autotexts:
            autotext.set_color('black')
            autotext.set_color('black')


        plt.tight_layout()

        # 保存路径
        bar_plot_path = opj(save_path, "bar_plot")
        md(bar_plot_path)

        plt.savefig(opj(bar_plot_path, "shap_summary_abs.pdf"), dpi=300, bbox_inches='tight')
        plt.close()
        print('shap_summary_abs已完成')

        # 6. 特征分布图（Beeswarm），创建 SHAP summary plot
        plt.figure()
        from matplotlib.colors import LinearSegmentedColormap

        # 创建 SHAP.Explanation 对象，手动指定 data 为全 0 或自定义数组，从而影响染色
        expl = shap.Explanation(
            values=shap_values,
            base_values=np.repeat(explainer.expected_value, X.shape[0]),
            data=X_plot.values,
            feature_names=X_plot.columns.tolist()
        )

        # 创建自定义 colormap：紫 → 橙
        # colors = ["#7a3db8", "#f3a54e"]
        # colors = ["#7a3db8", "#71DA56"]
        # colors = ["#B94FF7", "#71DA56"]

        # 使用 SHAP beeswarm 蜂群图（plot_type="dot"），颜色映射通过 expl.data 控制
        shap.plots.beeswarm(
            expl,
            show=False
        )

        # 获取坐标轴对象
        ax = plt.gca()

        # 1. 设置 y 轴特征标签字体
        ax.set_yticklabels(ax.get_yticklabels(), fontsize=14)

        # 2. 设置 x 轴刻度字体
        ax.tick_params(axis='x', labelsize=14)

        # 3. 设置 x/y 轴标题字体
        ax.set_xlabel("SHAP value (impact on model output)", fontsize=14)
        # ax.set_ylabel("Feature", fontsize=16)

        # 4. 修改散点大小
        for col in ax.collections:
            col.set_sizes([30])  # 默认约为 16，可设置为 30~60 比较合适
        fig = plt.gcf()

        # 设置 colorbar（颜色条）的刻度字体大小
        cbar = plt.gcf().axes[-1]  # 获取最后一个 axis（即 colorbar）
        cbar.tick_params(labelsize=14)  # 设置字体大小为 14
        # 找到 colorbar 的 label（即 "Feature value"）并设置字体大小
        cbar.set_ylabel("Feature value", fontsize=14, fontproperties=times_new_roman_prop)

        fig.set_size_inches(5.5, 4.5)
        beeswarm_plot_path = opj(save_path,"beeswarm_plot")
        md(beeswarm_plot_path)
        plt.savefig(opj(beeswarm_plot_path,"shap_summary.pdf"), dpi=300, bbox_inches='tight')
        plt.close()
        print('shap_summary已完成')

        y_proba = ensemble_predict_proba(X)
        # 在 non-ALNM 个体中，找概率最接近 0.421 的两个
        non_alnm_candidates = np.where(y == 0)[0]
        non_alnm_proba = y_proba[non_alnm_candidates]
        non_alnm_diff = np.abs(non_alnm_proba - 0.421)
        non_alnm_selected = non_alnm_candidates[np.argsort(non_alnm_diff)[:2]]

        # 在 ALNM 个体中，找概率最接近 0.421 的两个
        alnm_candidates = np.where(y == 1)[0]
        alnm_proba = y_proba[alnm_candidates]
        alnm_diff = np.abs(alnm_proba - 0.421)
        alnm_selected = alnm_candidates[np.argsort(alnm_diff)[:2]]
        # ============ 绘制1：bar_nonALNM_vs_ALNM_mean ============
        # 选择一位 non-ALNM 个体
        non_alnm_shap_1 = shap_values[non_alnm_selected[0]]
        non_alnm_shap_2 = shap_values[non_alnm_selected[1]]

        # 所有 ALNM 个体的平均 SHAP 值
        alnm_mean_shap = shap_values[np.where(y == 1)[0]].mean(axis=0)

        # 特征列表
        features = X_plot.columns.tolist()
        x = np.arange(len(features))
        bar_width = 0.25  # 三个柱子，所以宽度缩小

        # 绘图
        plt.figure(figsize=(8, 4))
        plt.bar(x - bar_width, non_alnm_shap_1, width=bar_width, color='#1E88E5', label='Non-ALNM case 1')
        plt.bar(x,            non_alnm_shap_2, width=bar_width, color='#2ca02c', label='Non-ALNM case 2')
        plt.bar(x + bar_width, alnm_mean_shap, width=bar_width, color='#FF0D57', label='Mean of ALNM cases')
        plt.axhline(0, color='black', linewidth=0.8)
        plt.xticks(x, features, fontsize=12)
        # 设置 y 轴刻度字体大小为 12
        plt.yticks(fontsize=12)
        plt.ylabel("SHAP value", fontsize=13)
        plt.legend(
            fontsize=12,
            loc='upper center',
            bbox_to_anchor=(0.5, -0.1),  # 图外底部居中
            frameon=False,
            ncol=3                        # 多列排版更紧凑
        )
        plt.tight_layout()
        plt.savefig(opj(bar_plot_path, "Non-ALNM_vs_ALNM_mean.pdf"), dpi=300, bbox_inches='tight')
        plt.close()
        print('tow-Non-ALNM_vs_ALNM_mean 已完成')

        # ============ 绘制3：bar_ALNM_vs_non-ALNM_mean ============
        # 选择一位 ALNM 个体
        alnm_shap_1 = shap_values[alnm_selected[0]]
        alnm_shap_2 = shap_values[alnm_selected[1]]
        # ============ 绘制2：bar_ALNM_vs_non-ALNM_mean ============
        # 选择一位 ALNM 个体
        alnm_shap_1 = shap_values[alnm_selected[0]]
        alnm_shap_2 = shap_values[alnm_selected[1]]
        # 所有 non-ALNM 个体的平均 SHAP 值
        non_alnm_mean_shap = shap_values[np.where(y == 0)[0]].mean(axis=0)

        # 特征列表
        features = X_plot.columns.tolist()
        x = np.arange(len(features))
        bar_width = 0.25

        # 绘图
        plt.figure(figsize=(8, 4))
        plt.bar(x - bar_width, alnm_shap_1, width=bar_width, color='#1E88E5', label='ALNM case 1')
        plt.bar(x,            alnm_shap_2, width=bar_width, color='#2ca02c', label='ALNM case 2')
        plt.bar(x + bar_width, non_alnm_mean_shap, width=bar_width, color='#FF0D57', label='Mean of Non-ALNM cases')
        plt.axhline(0, color='black', linewidth=0.8)
        plt.xticks(x, features, fontsize=12)
        # 设置 y 轴刻度字体大小为 12
        plt.yticks(fontsize=12)
        plt.ylabel("SHAP value", fontsize=13)
        plt.legend(
            fontsize=12,
            loc='upper center',
            bbox_to_anchor=(0.5, -0.1),  # 图外底部居中
            frameon=False,
            ncol=3                        # 多列排版更紧凑
        )
        plt.tight_layout()
        plt.savefig(opj(bar_plot_path, "ALNM_vs_Non-ALNM_mean.pdf"), dpi=300, bbox_inches='tight')
        plt.close()
        print('tow-ALNM_vs_Non-ALNM_mean 已完成')
        # ============================
        # 7. SHAP Dependence Plots（每个特征的 SHAP 值与原始值的关系）
        # ============================
        # 自定义 SHAP dependence plot，按 ALNM 分组显示颜色
        import seaborn as sns

        dependence_plot_path = opj(save_path, "dependence_plot")
        md(dependence_plot_path)
        # 设置颜色（0 类为紫色，1 类为橙色）
        # colors = {0: "#7a3db8", 1: "#f3a54e"}
        colors = {0: "#FF0D57", 1: "#1E88E5"}
        for feat, new_feat in zip(X.columns, X_plot.columns):
            fig, ax1 = plt.subplots(figsize=(4, 3.5))  # 主坐标轴（绘制 SHAP）
            ax2 = ax1.twinx()                          # 共享 x 轴（绘制直方图）

            # 1. 调整层级顺序
            ax2.set_zorder(0)  # 让直方图坐标轴在底层
            ax1.set_zorder(1)  # 让散点图坐标轴在顶层
            ax1.patch.set_visible(False)  # 把 ax1 的白色背景设为透明，避免遮挡直方图

            # 2. 绘制直方图（背景）
            ax2.hist(
                X[feat],
                bins=30,
                color='lightgrey',
                alpha=0.6,
                edgecolor='none'
            )
            ax2.set_yticks([])  # 不显示右侧 y 轴
            ax2.set_ylabel("Feature count", fontsize=12, color='gray')

            # 3. 绘制 SHAP 散点图（前景）
            for cls in [0, 1]:
                idx = y == cls
                ax1.scatter(
                    X.loc[idx, feat],
                    shap_values[idx, X.columns.get_loc(feat)],
                    alpha=0.9,
                    label="ALNM" if cls == 1 else "Non-ALNM",
                    color=colors[cls],
                    edgecolor="k",
                    linewidth=0,
                    s=30,
                    zorder=3  # 显式设置比直方图高
                )

            ax1.axhline(0, color='gray', linestyle='--', linewidth=1)

            # 4. 设置样式
            ax1.set_xlabel(new_feat, fontproperties=times_new_roman_prop, fontsize=14)
            ax1.set_ylabel("SHAP value", fontproperties=times_new_roman_prop, fontsize=14)
            ax1.tick_params(axis='x', labelsize=14)
            ax1.tick_params(axis='y', labelsize=14)
            ax1.legend(frameon=False, fontsize=14, handletextpad=0.02)

            plt.tight_layout()
            plt.savefig(opj(dependence_plot_path, f"{new_feat}.pdf"), dpi=300, bbox_inches='tight')
            plt.close()
        print('dependence_plot已完成')
    # ============================
    # 8. Waterfall Plot for a Single Sample
    # ============================

    waterfall_plot_path = opj(save_path, "waterfall")
    md(waterfall_plot_path)
    # 批量生成 Waterfall Plot
    for sample_index in range(N):
        sample_shap_values = shap_values[sample_index]
        sample_data = X_plot.iloc[sample_index]  # 替换列名后的数据

        expl = shap.Explanation(
            values=sample_shap_values,
            base_values=explainer.expected_value,
            data=sample_data,
            feature_names=X_plot.columns.tolist()
        )
        plt.figure(figsize=(6.5, 4))
        shap.plots.waterfall(expl, show=False)
        fig = plt.gcf()
        fig.set_size_inches(6, 4)
        label = y.iloc[sample_index]
        status = "ALNM" if label == 1 else "Non-ALNM"
        # === 获取预测概率
        sample_input = pd.DataFrame([X.iloc[sample_index]], columns=X.columns)
        prob = ensemble.predict_proba_by_dif(sample_input)[0][1]
        filename = opj(waterfall_plot_path,f"{status}_prob_{prob:.2f}_sample_{sample_index}.pdf")
        plt.tight_layout()
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
    print('Waterfall Plot已完成')
    # ============================
    # 9. SHAP Force Plot（力导向图）
    # ============================

    # # 保存为 HTML（推荐方式）
    force_plot_path = opj(save_path, "force_plot")
    md(force_plot_path)

    # Force Plot（保存 HTML）
    for i in range(N):
        force_plot = shap.force_plot(
            base_value=explainer.expected_value,
            shap_values=shap_values[i],
            features=X_plot.iloc[i],
            feature_names=X_plot.columns.tolist(),
            matplotlib=False
        )
        shap.save_html(opj(force_plot_path,f"sample_{i}.html"),force_plot)

    # # ============ 使用 shap.plots.force 生成堆叠图（HTML 可交互）============
    shap_html = shap.plots.force(
        base_value=explainer.expected_value,
        shap_values=shap_values[:N],
        features=X_plot.iloc[:N],
        feature_names=X_plot.columns.tolist(),
        matplotlib=False  # HTML 模式
    )
    shap.save_html(opj(force_plot_path, "a_stacked.html"), shap_html)
    
    print('叠加Force Plot已完成')